import numpy as np  # 导入NumPy库，用于数值计算
import pandas as pd  # 导入Pandas库，用于数据结构操作
import matplotlib.pyplot as plt  # 导入Matplotlib库，用于绘制图形
import time  # 导入time库，用于控制程序暂停时间

ALPHA = 0.1  # 学习率，决定每次Q值更新的幅度
GAMMA = 0.95  # 折扣因子，决定未来奖励的权重
EPSILION = 0.9  # epsilon-greedy策略中的探索概率，控制随机选择动作的比例
N_STATE = 20  # 状态空间的大小，表示状态的数量
ACTIONS = ['left', 'right']  # 可用的动作集合，左移和右移
MAX_EPISODES = 200  # 最大训练回合数
FRESH_TIME = 0.1  # 每步之间的时间间隔，用于渲染环境

# 构建Q表，Q表存储每个状态-动作对的Q值
def build_q_table(n_state, actions):
    q_table = pd.DataFrame(
        np.zeros((n_state, len(actions))),  # 创建一个形状为(n_state, len(actions))的全零矩阵
        np.arange(n_state),  # 状态的行索引为0到n_state-1
        actions  # 动作的列索引为'left'和'right'
    )
    return q_table  # 返回Q表

# 选择当前状态下的动作
def choose_action(state, q_table):
    # epsilon-greedy 策略
    state_action = q_table.loc[state, :]  # 获取当前状态下所有动作的Q值
    if np.random.uniform() > EPSILION or (state_action == 0).all():  # 探索（随机选择）或当Q值全为0时
        action_name = np.random.choice(ACTIONS)  # 随机选择一个动作
    else:  # 利用（选择Q值最大的动作）
        action_name = state_action.idxmax()  # 选择Q值最大对应的动作
    return action_name  # 返回选择的动作

# 获取环境反馈，依据当前状态和所选动作返回下一个状态和奖励
def get_env_feedback(state, action):
    if action == 'right':  # 如果选择了向右的动作
        if state == N_STATE - 2:  # 如果已经到达倒数第二个状态
            next_state = 'terminal'  # 终止状态
            reward = 1  # 到达终止状态时给予奖励1
        else:
            next_state = state + 1  # 否则，状态右移
            reward = -0.5  # 每步奖励为-0.5
    else:  # 如果选择了向左的动作
        if state == 0:  # 如果已经到达最左端
            next_state = 0  # 保持在状态0
        else:
            next_state = state - 1  # 否则，状态左移
        reward = -0.5  # 每步奖励为-0.5
    return next_state, reward  # 返回下一个状态和奖励

# 更新环境的状态并打印出来
def update_env(state, episode, step_counter):
    env = ['-'] * (N_STATE - 1) + ['T']  # 创建一个状态列表，其中T表示目标终止状态
    if state == 'terminal':  # 如果状态是终止状态
        print("Episode {}, the total step is {}".format(episode + 1, step_counter))  # 输出当前回合和步数
        final_env = ['-'] * (N_STATE - 1) + ['T']  # 终止状态时的环境状态
        return True, step_counter  # 返回终止状态和步数
    else:
        env[state] = '*'  # 将当前状态位置标记为*表示智能体所在的位置
        env = ''.join(env)  # 将状态列表转为字符串显示
        print(env)  # 打印当前环境状态
        time.sleep(FRESH_TIME)  # 暂停FRESH_TIME秒以控制显示速度
        return False, step_counter  # 返回非终止状态和步数

# Q-learning算法实现
def q_learning():
    q_table = build_q_table(N_STATE, ACTIONS)  # 构建Q表
    step_counter_times = []  # 用于存储每个回合的步数
    for episode in range(MAX_EPISODES):  # 遍历每个回合
        state = 0  # 每个回合从状态0开始
        is_terminal = False  # 是否到达终止状态的标志
        step_counter = 0  # 步数计数器
        update_env(state, episode, step_counter)  # 更新环境并显示
        while not is_terminal:  # 如果没有到达终止状态
            action = choose_action(state, q_table)  # 选择动作
            next_state, reward = get_env_feedback(state, action)  # 获取环境反馈
            next_q = q_table.loc[state, action]  # 获取当前状态-动作对的Q值
            if next_state == 'terminal':  # 如果到达终止状态
                is_terminal = True  # 标记为终止状态
                q_target = reward  # 目标Q值为奖励值
            else:  # 如果没有到达终止状态
                delta = reward + GAMMA * q_table.iloc[next_state, :].max() - q_table.loc[state, action]  # 计算TD误差
                q_table.loc[state, action] += ALPHA * delta  # 更新Q值
            state = next_state  # 更新状态
            is_terminal, steps = update_env(state, episode, step_counter + 1)  # 更新环境并显示
            step_counter += 1  # 步数+1
            if is_terminal:  # 如果到达终止状态
                step_counter_times.append(steps)  # 记录回合的步数
    return q_table, step_counter_times  # 返回更新后的Q表和每个回合的步数列表


# 程序入口
if __name__ == '__main__':
    q_table, step_counter_times = q_learning()  # 运行Q-learning算法
    print("Q table\n{}\n".format(q_table))  # 打印最终的Q表
    print('end')  # 打印训练结束信息

    # 绘制每回合步数的图表
    plt.plot(step_counter_times, 'g-')  # 以绿色线条绘制步数
    plt.ylabel("steps")  # 设置Y轴标签为"steps"
    plt.title("Q-Learning Algorithm")  # 设置图标题
    plt.show()  # 显示图表
    print("The step_counter_times is {}".format(step_counter_times))  # 打印每个回合的步数
